#include <stdio.h>
#include <cuda_runtime.h>

#define checkCudaErrors(err) __checkCudaErrors(err, __FILE__, __LINE__)

inline void __checkCudaErrors(cudaError_t err, const char *file, const int line) {
	if (err != cudaSuccess) {
		printf("CUDA Error: %s in %s at line %d\n", cudaGetErrorString(err), file, line);
		exit(EXIT_FAILURE);
	}
}

__global__ void add(const int *a, const int *b, int *c, int n) {
	int index = threadIdx.x + blockIdx.x * blockDim.x;
	if (index < n) {
		c[index] = a[index] + b[index];
	}
}

int main() {
	int n = 1024; // 向量大小
	size_t bytes = n * sizeof(int);
	unsigned int flags;
	int *h_a, *h_b, *h_c; // 主机端向量
	int *d_a, *d_b, *d_c; // 设备端向量
	cudaStream_t stream; // CUDA流

	// check env
	int idev = 0; // use default device 0
	cudaDeviceProp deviceProp;

	printf("CUDART_VERSION = %d\n", CUDART_VERSION);
	if (CUDART_VERSION < 2020) {
		printf("CUDART_VERSION < 2020, please upgrade your CUDA driver\n");
		return 1;
	}

	cudaSetDevice(idev);
	cudaGetDeviceProperties(&deviceProp, idev);

	// need: CUDART_VERSION >= 2020
	if (!deviceProp.canMapHostMemory) {
		fprintf(stderr, "Device %d does not support mapping CPU host memory!\n",
				idev);

		exit(EXIT_SUCCESS);
	}

	cudaSetDeviceFlags(cudaDeviceMapHost);

	// 分配主机内存
	flags = cudaHostAllocMapped;
	checkCudaErrors(cudaHostAlloc((void **)&h_a, bytes, flags));
	checkCudaErrors(cudaHostAlloc((void **)&h_b, bytes, flags));
	checkCudaErrors(cudaHostAlloc((void **)&h_c, bytes, flags));

	// 初始化主机端向量
	for (int i = 0; i < n; i++) {
		h_a[i] = i;
		h_b[i] = 2 * i;
	}

	// 获取设备指针
	checkCudaErrors(cudaHostGetDevicePointer((void **)&d_a, (void *)h_a, 0));
	checkCudaErrors(cudaHostGetDevicePointer((void **)&d_b, (void *)h_b, 0));
	checkCudaErrors(cudaHostGetDevicePointer((void **)&d_c, (void *)h_c, 0));

	// 创建流
	checkCudaErrors(cudaStreamCreate(&stream));

	// 由于数据已经在锁页内存中，我们可以直接在 CUDA 内核中使用它
	// 执行内核
	int threadsPerBlock = 256;
	int blocksPerGrid = (n + threadsPerBlock - 1) / threadsPerBlock;
	add<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(d_a, d_b, d_c, n);

	// 等待流中的所有操作完成
	checkCudaErrors(cudaStreamSynchronize(stream));

	// 检查结果
	int i;
	for (i = 0; i < n; i++) {
		if (h_c[i] != h_a[i] + h_b[i]) {
			printf("Error: invalid result at %d\n", i);
			exit(EXIT_FAILURE);
		}
	}

	if (i >= n) {
		printf("The result is correct\n");
	}

	// 清理资源
	cudaFreeHost(h_a);
	cudaFreeHost(h_b);
	cudaFreeHost(h_c);
	cudaStreamDestroy(stream);

	return 0;
}
